import torch
import albumentations as A                         # 图像增强
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm                              # 进度条提示
import torch.nn as nn
from model import Unet
import torch.optim as optim
from utils import (
get_loaders,
check_accuracy,
save_predictions_as_image,
)

# 超参数部分
LEARNING_RATE = 1E-5
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
BATCH_SIZE = 16
NUM_EPOCHS = 1
NUM_WORKERS = 5    #线程
IMAGE_HEIGHT = 160
IMAGE_WIDTH = 240
LOAD_MODEL = False    # 是否加载预训练模型
TRAIN_IMG_DIR='./data/images/'
TRAIN_MASK_DIR='./data/masks/'
VAL_IMG_DIR='./data/val_images/'
VAL_MASK_DIR = './data/val_masks/'

# 一个epoch的训练
def train_fn(loader,model,optimizer,loss_fn,scaler):
    loop=tqdm(loader)
    for img, label in loop:
        img=img.to(device=DEVICE)
        label=label.float().unsqueeze(1).to(DEVICE)

        # forward
        with torch.cuda.amp.autocast():
            predictions = model(img)
            loss = loss_fn(predictions,label)

        #backward
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        # update
        loop.set_postfix(loss=loss.item())

def main():
    #训练集预处理
    train_transforms=A.Compose(
        [
            A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
            A.Rotate(limit=35,p=0.5),     #-35~35的旋转，旋转概率0.5
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.1),

            A.Normalize(
                mean=[0.0,0.0,0.0],
                std=[1.0,1.0,1.0],
                max_pixel_value=255.0
            ),
            ToTensorV2(),
        ]
    )

    #验证集数据预处理
    val_transforms=A.Compose(
        [
            A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
            A.Normalize(
                mean=[0.0,0.0,0.0],
                std=[1.0,1.0,1.0],
                max_pixel_value=255.0
            ),
            ToTensorV2(),
        ]
    )

    #实例化UNET、loss、optimaizer
    model=Unet(inChannel=3,outChannel=1).to(DEVICE)
    loss_fn=nn.BCEWithLogitsLoss()
    optimizer =optim.Adam(model.parameters(), lr=LEARNING_RATE)

    #获取数据集
    train_loader, val_loader = get_loaders(
        TRAIN_IMG_DIR,
        TRAIN_MASK_DIR,
        VAL_IMG_DIR,
        VAL_MASK_DIR,
        BATCH_SIZE,
        train_transforms,
        val_transforms,
        NUM_WORKERS,
    )

    # 加载预训练权重
    if LOAD_MODEL:
        print('Pretrained:')
        model.load_state_dict(torch.load('unet.pth'))
        check_accuracy(val_loader, model, device=DEVICE)
        print('------>Loading pretrained model successfully!!')

    scaler=torch.cuda.amp.GradScaler()

    for epoch in range(NUM_EPOCHS):
        print('Epoch:', epoch + 1)
        train_fn(train_loader, model, optimizer, loss_fn, scaler)  # 训练一个 epoch

        # check accuracy
        check_accuracy(val_loader, model, device=DEVICE)

        # save model
    print('------>Saving checkpoint')
    torch.save(model.state_dict(), 'unet.pth')

    # print some examples to a folder
    save_predictions_as_image(val_loader, model, folder='saved_val_images/', device=DEVICE)


if __name__ == '__main__':
    main()
    print(' training over!!!! ')
